'''
Description: play motion
Version: V1.0
Author: zw_1520@163.com
Date: 2024-11-19 15:51:41
LastEditors: zw_1520@163.com
LastEditTime: 2025-05-06 15:47:44
Copyright (C) All rights reserved.
'''

import argparse
import signal
import sys
import threading
import time

import aimdk.protocol_aimrt_rpc_pb2
import aimdk.protocol_pb2
import aimrt_py
import yaml

global_aimrt_core = None
running = threading.Event()
running.set()


def signal_handler(sig, _):
    global global_aimrt_core, running

    if (global_aimrt_core and (sig == signal.SIGINT or sig == signal.SIGTERM)):
        running.clear()
        global_aimrt_core.Shutdown()
        return

    sys.exit(0)


def play_motion(module_logger, mc_action_proxy, motion_cmd_proxy, motion_id, duration_ms):
    # Gets current action status
    req = aimdk.protocol_pb2.McActionRequest()
    status, rsp = mc_action_proxy.GetAction(req)
    if not status.OK():
        aimrt_py.error(
            module_logger, f"Get action failed. status: {status.ToString()}")
        return
    aimrt_py.info(module_logger, f"Get action success. action: {rsp}")

    VALID_ACTIONS = (
        aimdk.protocol_pb2.McAction.McAction_STAND_DEFAULT,
        aimdk.protocol_pb2.McAction.McAction_RL_LOCOMOTION_ARM_EXT_JOINT_SERVO
    )
    if rsp.info.current_action not in VALID_ACTIONS:
        aimrt_py.error(
            module_logger, f"Current action is not valid. action: {rsp.info.current_action}")
        # return

    # Changes to motion action
    req = aimdk.protocol_pb2.McActionRequest()
    req.command.action = aimdk.protocol_pb2.McAction.McAction_RL_LOCOMOTION_ARM_EXT_JOINT_SERVO
    status, rsp = mc_action_proxy.SetAction(req)

    if not status.OK():
        aimrt_py.error(
            module_logger, f"Set action failed. status: {status.ToString()}")
        return

    aimrt_py.info(
        module_logger, f"Set action to STAND_ARM_EXT_JOINT_SERVO success.")

    # Plays motion
    req = aimdk.protocol_pb2.MotionCommandRequest()
    req.motion_id = motion_id
    req.duration_ms = duration_ms
    req.cmd_end = True
    req.cmd_pause = False
    req.cmd_reset = False
    status, rsp = motion_cmd_proxy.SendMotionCommand(req)
    if not status.OK():
        aimrt_py.error(
            module_logger, f"Play motion failed. status: {status.ToString()}")
        return

    aimrt_py.info(
        module_logger, f"Play motion success. motion_id: {motion_id}, duration_ms: {duration_ms}")


def main():
    parser = argparse.ArgumentParser(description='Example rpc server app.')
    parser.add_argument('--cfg_file_path', type=str,
                        default="", help='config file path')
    args = parser.parse_args()

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    print("AimRT start.")

    aimrt_core = aimrt_py.Core()

    global global_aimrt_core
    global_aimrt_core = aimrt_core

    # Initializes
    core_options = aimrt_py.CoreOptions()
    core_options.cfg_file_path = args.cfg_file_path
    aimrt_core.Initialize(core_options)

    # Creates module
    module_handle = aimrt_core.CreateModule("PlayMotionModule")

    # Reads cfg
    module_cfg_file_path = module_handle.GetConfigurator().GetConfigFilePath()
    with open(module_cfg_file_path, 'r') as file:
        data = yaml.safe_load(file)
        motion_id = data["motion_id"]
        duration_ms = int(data["duration_ms"])

    # Registers rpc client
    rpc_handle = module_handle.GetRpcHandle()
    mc_action_proxy = aimdk.protocol_aimrt_rpc_pb2.McActionServiceProxy(
        rpc_handle)
    ret = mc_action_proxy.RegisterClientFunc(rpc_handle)
    assert ret, "Register pb rpc service failed."

    motion_cmd_proxy = aimdk.protocol_aimrt_rpc_pb2.MotionCommandServiceProxy(
        rpc_handle)
    ret = motion_cmd_proxy.RegisterClientFunc(rpc_handle)
    assert ret, "Register pb rpc service failed."

    # Starts
    thread = threading.Thread(target=aimrt_core.Start)
    thread.start()

    # Sleeps
    time.sleep(1)

    # Plays motion
    play_motion(module_handle.GetLogger(), mc_action_proxy,
                motion_cmd_proxy, motion_id, duration_ms)

    # Exits
    aimrt_core.Shutdown()

    while thread.is_alive():
        thread.join(1.0)

    print("AimRT exit.")


def play_motion_by_external(config_file_path, motion_id, duration_ms):
    aimrt_core = aimrt_py.Core()

    global global_aimrt_core
    global_aimrt_core = aimrt_core

    # Initializes
    core_options = aimrt_py.CoreOptions()
    core_options.cfg_file_path = config_file_path
    aimrt_core.Initialize(core_options)

    # Creates module
    module_handle = aimrt_core.CreateModule("PlayMotionModule")

    # Reads cfg
    module_cfg_file_path = module_handle.GetConfigurator().GetConfigFilePath()
    with open(module_cfg_file_path, 'r') as file:
        data = yaml.safe_load(file)
        # motion_id = data["motion_id"]
        # duration_ms = int(data["duration_ms"])

    # Registers rpc client
    rpc_handle = module_handle.GetRpcHandle()
    mc_action_proxy = aimdk.protocol_aimrt_rpc_pb2.McActionServiceProxy(
        rpc_handle)
    ret = mc_action_proxy.RegisterClientFunc(rpc_handle)
    assert ret, "Register pb rpc service failed."

    motion_cmd_proxy = aimdk.protocol_aimrt_rpc_pb2.MotionCommandServiceProxy(
        rpc_handle)
    ret = motion_cmd_proxy.RegisterClientFunc(rpc_handle)
    assert ret, "Register pb rpc service failed."

    # Starts
    thread = threading.Thread(target=aimrt_core.Start)
    thread.start()

    # Sleeps
    time.sleep(1)

    # Plays motion
    print(motion_id)
    print(duration_ms)
    play_motion(module_handle.GetLogger(), mc_action_proxy,
                motion_cmd_proxy, motion_id, duration_ms)

    # Exits
    aimrt_core.Shutdown()

    while thread.is_alive():
        thread.join(1.0)

    print("AimRT exitxxx.")


if __name__ == '__main__':
    main()
